import datetime
import pandas as pd
from gofishing.marketData import _pvFromSQLiteDB
from gofishing.base import TradeType, BaseBacktest, Log
from gofishing.commodity import CommodityBacktest
from gofishing.indicator import ATR, Donchian
from gofishing.strategy import Strategy
from gofishing.observer import AccountObserver
from gofishing.fishing import Fishing
from gofishing.analyse import Analyzer


class StrategyTest(Strategy):

    def __init__(self, stratID: str, commodity: CommodityBacktest,
                 donchianLong: Donchian, donchianShort: Donchian,
                 atr: ATR) -> None:
        super().__init__(stratID=stratID, commodity=commodity)
        self.donchianLong = donchianLong
        self.donchianShort = donchianShort
        self.atr = atr
        self.maxLong = 1
        self.volumeEvery = 1

    def updateDecision(self) -> None:
        dcLongHigh = self.donchianLong.high
        dcShortLow = self.donchianShort.low
        atr = self.atr.atr
        high = self.commodity.high
        low = self.commodity.low
        close = self.commodity.close

        if self.long + self.longOpening < self.maxLong and \
                high[-2] < dcLongHigh[-2] and close[-1] >= dcLongHigh[-2]:
            openPrice = close[-1]
            self.open(tradeType=TradeType.LONGOPEN,
                      limitPrice=openPrice, volume=self.volumeEvery)

        canClose = self.long - self.longClosing
        if canClose > 0 and \
                low[-2] > dcShortLow[-2] and close[-1] <= dcShortLow[-2]:
            self.close(tradeType=TradeType.LONGCLOSE,
                       limitPrice=close[-1], volume=canClose)
        elif canClose > 0 and \
                close[-1] < self.longOpenPrice - 2 * atr[-1]:
            self.close(tradeType=TradeType.LONGCLOSE,
                       limitPrice=close[-1], volume=canClose)

    def onCloseOfDay(self) -> None:
        # 切换主力合约
        if not self.commodity.isDominant():
            if self.isNoPosition():
                self.commodity.switchToDominant()


def runBacktest() -> None:
    codes = [
        'A2101.XDCE', 'A2103.XDCE', 'A2105.XDCE', 'A2107.XDCE',
        'A2109.XDCE', 'A2111.XDCE',
        'A2201.XDCE', 'A2203.XDCE', 'A2205.XDCE', 'A2207.XDCE',
        'A2209.XDCE', 'A2211.XDCE']

    start = '2020-10-1'
    end = '2021-10-15'

    BaseBacktest.init(
        _pvFromSQLiteDB(
            'A8888.XDCE', '5m', start=start, end=end).index)
    Log.init(logName='backtest')

    recordIndex = \
        _pvFromSQLiteDB('A8888.XDCE', '1d').index + pd.Timedelta('15h')
    AccountObserver(observerID='accountObserver', recordIndex=recordIndex)

    commodityA = CommodityBacktest(commodityID='A',
                                    contractCodes=codes,
                                    frequency='15m',
                                    frequencyHigh='5m',
                                    length=25)
    StrategyTest(
        stratID='strategyTest',
        commodity=commodityA,
        donchianLong=Donchian(indicatorID='donchianLong',
                                commodity=commodityA,
                                period=20),
        donchianShort=Donchian(indicatorID='donchianShort',
                                commodity=commodityA,
                                period=10),
        atr=ATR(indicatorID='ATR', commodity=commodityA))

    # 开始回测
    Fishing.start(board=False)
    Analyzer.writeRecord(prefix='backtest')
    Analyzer.printBrokerBrief()


if __name__ == '__main__':
    t0 = datetime.datetime.now()

    runBacktest()

    t = datetime.datetime.now()
    print('Time consumed: %s' % (t-t0))
